import os
from openai import OpenAI

# 请确保您已将 API Key 存储在环境变量 ARK_API_KEY 中
# 初始化Openai客户端，从环境变量中读取您的API Key
client = OpenAI(
    # 此为默认路径，您可根据业务所在地域进行配置
    base_url="https://ark.cn-beijing.volces.com/api/v3",
    # 从环境变量中获取您的 API Key
    api_key='d7138b36-834d-4fb1-836a-b74b103d409b',
)

"""
deepseek-v3-250324:ep-20250327145610-kbll5
deepseek-v3:ep-20250310173311-wgx5m
deepseek-R1:ep-20250212112426-prrlr
deepseek-R1-250528:ep-20250609120241-jkn4v

"""
def get_standard_response(messages, model="ep-20250327145610-kbll5"):
    """
    获取标准（非流式）响应
    :param messages: 消息列表
    :param model: 模型ID
    :return: 响应内容
    """
    completion = client.chat.completions.create(
        model=model,
        messages=messages,
        max_tokens=1024,  # 生成的最大token数
        top_p=0.5,  # 控制输出文本的多样性
        frequency_penalty=0.0,  # 控制生成文本中重复单词的频率
        presence_penalty=0.0,  # 控制生成文本中新的主题或词汇的引入
        n=1,  # 返回的响应数量
    )
    return completion.choices[0].message.content

def get_streaming_response(messages, model="ep-20250327145610-kbll5"):
    """
    获取流式响应
    :param messages: 消息列表
    :param model: 模型ID
    :return: 生成器，用于逐步获取响应内容
    """
    stream = client.chat.completions.create(
        model=model,
        messages=messages,
        stream=True,
        max_tokens=1024,  # 生成的最大token数
        top_p=0.5,  # 控制输出文本的多样性
        frequency_penalty=0.0,  # 控制生成文本中重复单词的频率
        presence_penalty=0.0,  # 控制生成文本中新的主题或词汇的引入
        n=1,  # 返回的响应数量
    )
    
    for chunk in stream:
        if chunk.choices:
            yield chunk.choices[0].delta.content or ""


def get_standard_response_R1(messages, model="ep-20250609120241-jkn4v"):
    """
    获取标准（非流式）响应
    :param messages: 消息列表
    :param model: 模型ID
    :return: 响应内容
    """
    completion = client.chat.completions.create(
        model=model,
        messages=messages,
        max_tokens=1024,  # 生成的最大token数
        top_p=0.5,  # 控制输出文本的多样性
        frequency_penalty=0.0,  # 控制生成文本中重复单词的频率
        presence_penalty=0.0,  # 控制生成文本中新的主题或词汇的引入
        n=1,  # 返回的响应数量
    )
    return completion.choices[0].message.content

def get_streaming_response_R1(messages, model="ep-20250609120241-jkn4v"):
    """
    获取流式响应
    :param messages: 消息列表
    :param model: 模型ID
    :return: 生成器，用于逐步获取响应内容
    """
    stream = client.chat.completions.create(
        model=model,
        messages=messages,
        stream=True,
        max_tokens=1024,  # 生成的最大token数
        top_p=0.5,  # 控制输出文本的多样性
        frequency_penalty=0.0,  # 控制生成文本中重复单词的频率
        presence_penalty=0.0,  # 控制生成文本中新的主题或词汇的引入
        n=1,  # 返回的响应数量
    )
    
    for chunk in stream:
        if chunk.choices:
            yield chunk.choices[0].delta.content or ""



# 使用示例
if __name__ == "__main__":
    test_messages = [
        {"role": "system", "content": "你是人工智能助手"},
        {"role": "user", "content": "你好，你是？"},
    ]
    
    # 测试标准响应
    # print("----- standard request -----")
    # response = get_standard_response(test_messages)
    # print(response)
    
    # 测试流式响应
    print("----- streaming request -----")
    for content in get_streaming_response_R1(test_messages):
        print(content, end="")
    print()
